{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    }
   },
   "source": [
    "# Solving Inverted Pendulum with SVP\n",
    "\n",
    "This notebook shows how to use the proposed structured value-based planning (SVP) approach to generate the state-action $Q$-value function for the classic inverted pendulum problem. The correctness of the solution is verified by trajectory simulations."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    }
   },
   "source": [
    "## Problem definition\n",
    "\n",
    "The goal is to balance the inverted pendulum on the upright equilibrium position.\n",
    "The physical dynamics of the system is described by the angle and the angular speed, i.e., $(\\theta, \\dot{\\theta})$. Denote $\\tau$ as the time interval between decisions, $u$ as the torque input on the pendulum, the dynamics can be written as\n",
    "\n",
    "\\begin{align}\n",
    "    & \\theta := \\theta + \\dot{\\theta}~\\tau,\\\\\n",
    "    & \\dot{\\theta} := \\dot{\\theta} + \\left( \\sin{\\theta} - \\dot{\\theta} + u \\right)\\tau.\n",
    "\\end{align}\n",
    "\n",
    "A reward function that penalizes control effort while favoring an upright pendulum is used as\n",
    "\n",
    "\\begin{equation}\n",
    "    r(\\theta,u) = - 0.1u^2 + \\exp{\\left(\\cos{\\theta}-1 \\right)}.\n",
    "\\end{equation}\n",
    "\n",
    "In the simulation, the state space is $(-\\pi, \\pi]$ for $\\theta$ and $[-10,10]$ for $\\dot{\\theta}$. We limit the input torque in $[-1,1]$ and set $\\tau=0.3$.\n",
    "We discretize each dimension of the state space into 50 values, and action space into 1000 values, which forms an $Q$-value function  a matrix of dimension $2500\\times 1000$."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    }
   },
   "source": [
    "## Structured Value-based Planning (SVP)\n",
    "\n",
    "The proposed structured value-based planning (SVP) approach is based on the $Q$-value iteration. At the $t$-th iteration, instead of a full pass over all state-action pairs:\n",
    "- SVP first randomly selects a subset $\\Omega$ of the state-action pairs. In particular, each state-action pair in $\\mathcal{S}\\times\\mathcal{A}$ is observed (i.e., included in $\\Omega$) independently with probability $p$. \n",
    "- For each selected $(s,a)$, the intermediate $\\hat{Q}(s,a)$ is computed based on the $Q$-value iteration: \n",
    "    \\begin{equation*}\\hat{Q}(s,a) \\leftarrow \\sum_{s'} P(s'|s,a) \\left( r(s,a) + \\gamma \\max_{a'} Q^{(t)}(s',a') \\right),\\quad\\forall\\:(s,a)\\in\\Omega.\n",
    "    \t\t\t\t\\end{equation*}\n",
    "- The current iteration then ends by reconstructing the full $Q$ matrix with matrix estimation, from the set of observations in $\\Omega$. That is, $Q^{(t+1)}=\\textrm{ME}\\big(\\{\\hat{Q}(s,a)\\}_{(s,a)\\in\\Omega}\\big).$\n",
    "\n",
    "Overall, each iteration reduces the computation cost by roughly $1-p$."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "subslide"
    }
   },
   "source": [
    "Through SVP, we can compute the final state-action $Q$-value function.\n",
    "To obtain the optimal policy for state $s$, we compute\n",
    "\n",
    "\\begin{align*}\n",
    "    \\pi^{\\star} \\left(s\\right) = \\mbox{argmax}_{a \\in \\mathcal{A}} Q^{\\star}\\left(s, a\\right).\n",
    "\\end{align*}"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    }
   },
   "source": [
    "### Generate state-action value function with SVP"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "slideshow": {
     "slide_type": "fragment"
    }
   },
   "outputs": [],
   "source": [
    "push!(LOAD_PATH, \".\")\n",
    "using MDPs, InvertedPendulum, Printf, LinearAlgebra\n",
    "mdp = MDP(state_space(), action_space(), transition, reward)\n",
    "__init__()\n",
    "policy = value_iteration(mdp, true, \"../data/qip_otf_0.4.csv\", true)\n",
    "print(\"\")  # suppress output"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    }
   },
   "source": [
    "### Visualize policy as a heat map"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "slideshow": {
     "slide_type": "fragment"
    }
   },
   "outputs": [],
   "source": [
    "viz_policy(mdp, policy, \"SVP policy (40% observed)\", true, \"ip/policy_ip_0.4.tex\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    }
   },
   "source": [
    "### Verify correctness\n",
    "Simulate and visualize trajectory from initial state `[angle, speed]`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "slideshow": {
     "slide_type": "fragment"
    }
   },
   "outputs": [],
   "source": [
    "ss, as = simulate(mdp, policy, [-0.5, 0.0])\n",
    "viz_trajectory(ss, as, \"SVP policy trajectory (40% observed)\", \"SVP policy input (40% observed)\", true, \"ip/traj_ip_0.4.tex\")\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "nsim = 2000\n",
    "deviation = 0\n",
    "for sim = 1:nsim\n",
    "    state = [rand(PMIN:0.001 * (PMAX - PMIN):PMAX), \n",
    "             rand(VMIN:0.001 * (VMAX - VMIN):VMAX)]\n",
    "    traj, _ = simulate(mdp, policy, copy(state))\n",
    "    deviation += vecnorm(traj[51:end, 1])\n",
    "end # for sim\n",
    "deviation /= nsim\n",
    "@printf(\"average deviation: %.3f\\n\", deviation)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "celltoolbar": "Slideshow",
  "kernelspec": {
   "display_name": "Julia 0.7.0",
   "language": "julia",
   "name": "julia-0.7"
  },
  "language_info": {
   "file_extension": ".jl",
   "mimetype": "application/julia",
   "name": "julia",
   "version": "0.7.0"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 1
}
